//
// Created by yanhai on 2021/1/7.
//

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>

#define RED     1
#define BLACK   2

typedef int KEY_TYPE;

typedef struct rbtree_node_ {
    unsigned char color;
    struct rbtree_node_ *left;
    struct rbtree_node_ *right;
    struct rbtree_node_ *parent;
    KEY_TYPE key;
    void *value;    // value
} rbtree_node;

typedef struct rbtree_ {
    rbtree_node *root;
    rbtree_node *nil;
} rbtree;

/**
 *           M(y)                   P(x)
 *         /  \        右旋        /  \
 *       P(x) 子树2    ----->    子树1  M(y)
 *     /  \           <-----         /  \
 *  子树1   N           左旋         N    子树2
 *
 * 经过上图可以发现，左旋和右旋 就是 P M N 3个节点之间的相互关系发生了变化
 * 因此，旋转就是修改 P M N 节点的6个指针的方向
 *
 * 以对P进行左旋为例，P节点为x
 * 1. 保存 x->right 为y，即M节点
 * 2. 首先建立P与N的关系，N有可能为是叶子节点，因此需要判断一下
 *      x->right = y->left;
 *      y->left->parent = x;
 * 3. 建立M和P的父节点的关系，
 *      y->parent = x->parent;
 *      根据P的位置（是根节点，还是在左子树，还是在右子树）修改它的父节点指向M节点
 *      1. 如果 P是根节点，或者判断P的父节点为nil，那么M节点就变为根节点
 *          T->root = y;
 *      2. 如果P在左子树上
 *          x->parent->left = y;
 *      3. 如果P在右子树上
 *          x->parent->right = y;
 * 4. 最后建立M和P的关系
 *      y->left = x;
 *      x->parent = y;
 */

/**
 * 左旋
 * 左旋就是本身节点向左下沉，右子节点上升到本身节点位置，使本身节点成为原本它右子节点的左子节点
 * 左旋是将某个节点旋转为其右孩子的左孩子
 * @param T
 * @param x
 */
void rbtree_left_rotate(rbtree *T, rbtree_node *x) {
    // x就是右图中的P节点，y就是右图中的M节点，y->left就是N节点
    // y一定不是nil，因此不用判断
    rbtree_node *y = x->right;

    // y->left就是右图中的N节点
    // P节点和N节点建立关系
    // P -> N
    x->right = y->left;         // 修改了第1个指针的方向
    if (y->left != T->nil) {
        // N -> P
        y->left->parent = x;    // 修改了第2个指针的方向
    }

    // M和P的父节点建立关系
    y->parent = x->parent;      // 3
    if (x->parent == T->nil) {  // 表示x就是根节点
        T->root = y;            // 4
    } else if (x == x->parent->left) {    // P在左子树上面
        x->parent->left = y;
    } else {                    // P在右子树上面
        x->parent->right = y;
    }

    // M和P建立连接
    y->left = x;                // 5
    x->parent = y;              // 6
}

/**
 * 右旋
 * 右旋就是本身节点向右下沉，左子节点上升到本身节点位置，使本身节点成为原本它左子节点的右子节点
 * 右旋是节点旋转为其左孩子的右孩子
 * @param T
 * @param y
 */
void rbtree_right_rotate(rbtree *T, rbtree_node *y) {
    rbtree_node *x = y->left;

    y->left = x->right;
    if (x->right != T->nil) {
        x->right->parent = y;
    }

    x->parent = y->parent;
    if (y->parent == T->nil) {
        T->root = x;
    } else if (y == y->parent->right) {
        y->parent->right = x;
    } else {
        y->parent->left = x;
    }

    x->right = y;
    y->parent = x;
}

/**
 * 插入后，有5种情况，其中2种不需要调整其他节点，3种需要调整其他节点
 * 前2种：
 *  1. 插入的节点为根节点：只需要改为黑色即可。
 *  2. 插入的节点的父节点为黑色：什么都不需要改。
 * 剩下3中均需要调整其他节点
 * 注意：插入后需要调整时的前提条件
 *  1. 本身节点为红色
 *  2. 父节点为红色
 * 调整的3种情况：
 *  1. 叔父节点为红色
 *  2. 叔父节点为黑色，并且本身节点在左子树上
 *  3. 叔父节点为黑色，并且本身节点在右子树上
 */
void rbtree_insert_fixup(rbtree *T, rbtree_node *z) {
    // 由于我们新插入的节点的颜色是红色，
    // 因此，如果它的父节点是黑色的，那就不需要调整了，
    // 如果它的父节点是红色的，那就需要继续调整
    while (z->parent->color == RED) {
        // 如果z的父节点在左子树
        // todo 这里无需担心z->parent->parent为nil，
        //      因为z->parent是红色的，z->parent->parent一定存在，并且为黑色
        if (z->parent == z->parent->parent->left) {
            // y是z父节点的兄弟节点（z的叔父节点）
            rbtree_node *y = z->parent->parent->right;
            // z节点、z的父节点和叔父节点同为红色，
            if (y->color == RED) {
                // 现在只需要把父节点和叔父节点改为黑色，祖父节点改为红色即可
                // 此时这3层就是平衡的了，我们只需要继续判断由于祖父节点的改变，对上层节点的影响即可。
                z->parent->color = BLACK;   // 父亲节点
                y->color = BLACK;       // 叔父节点
                z->parent->parent->color = RED; // 祖父节点

                // 由于祖父节点被改成了红色，如果他的父亲节点也是红色，
                // 则仍然需要调整，将z指向祖父节点，进行递归
                z = z->parent->parent;
            } else {    // z节点、z的父节点为红色， 叔父节点为黑色
                // 如果z在右子树上面，需要对z->parent向左旋转
                // 使z->parent成为z的左节点
                if (z == z->parent->right) {
                    z = z->parent;  // 修改z的指向，指向它的父节点，父节点经过左旋，父节点变为左子树节点
                    rbtree_left_rotate(T, z);
                }
                // z->parent就是上的z，新插入的那个节点
                z->parent->color = BLACK;
                z->parent->parent->color = RED;
                rbtree_right_rotate(T, z->parent->parent);  // 此时z为红色，z->parent为黑色，结束循环
            }
        } else { // z的父节点在右子树（理解上面的就可以了，原理是一样的）
            // y是z的叔父节点
            rbtree_node *y = z->parent->parent->left;
            if (y->color == RED) {
                z->parent->color = BLACK;
                y->color = BLACK;
                z->parent->parent->color = RED;

                z = z->parent->parent;
            } else {
                if (z == z->parent->left) {
                    z = z->parent;
                    rbtree_right_rotate(T, z);
                }
                z->parent->color = BLACK;
                z->parent->parent->color = RED;
                rbtree_left_rotate(T, z->parent->parent);
            }
        }
    }

    // TODO 为什么每次都需要设置一下根节点的颜色，在上面的循环中没有设置根节点的颜色吗？
    // 在调整的过程中，有可能会把根节点的颜色改为红色，因此在这里用重新改回来
    // 比如经过了上面的调整，第2层2个节点都是黑色的，那么根节点就被改为红色了
    // 又比如，插入的节点是根节点
    T->root->color = BLACK;
}

void rbtree_insert(rbtree *T, rbtree_node *z) {
    // 这里不可以省略y变量，因为x为叶子节点时，
    // 它虽然有x->parent，但是x->parent却是不确定的
    // 我们从来没有对T->nil->parent赋过值

    // 先找到z要插入的位置，和二叉树的insert是一样的
    rbtree_node *x = T->root;
    rbtree_node *y = T->nil;
    while (x != T->nil) {
        y = x;
        if (z->key < x->key) {
            x = x->left;
        } else if (z->key > x->key) {
            x = x->right;
        } else {    // exist
            return;
        }
    }

    // 此时叶子节点x的位置就是存放z的位置
    // 它对于的父节点就是y，因此z的父节点是y
    z->parent = y;
    // 下面建立y->z的关系
    if (y == T->nil) {
        T->root = z;
    } else if (z->key < y->key) {
        y->left = z;
    } else {
        y->right = z;
    }

    z->left = T->nil;
    z->right = T->nil;
    // TODO 为什么新插入的要设置为红色？
    // 因为如果插入的是黑色，会影响黑高，每次都需要调整，并且调整起来比较麻烦
    z->color = RED;

    // 下面就是关键了，插入节点有可能会影响红黑树的规则，因此要调整
    rbtree_insert_fixup(T, z);
}

// todo ...
int rbtree_delete(rbtree *T, rbtree_node *z) {

    return 0;
}

void rbtree_traversal(rbtree *T, rbtree_node *node) {
    if (node != T->nil) {
        rbtree_traversal(T, node->left);
        printf("%4d", node->key);
        rbtree_traversal(T, node->right);
    }
}

void rbtree_traversal_wrapper(rbtree *T) {
    rbtree_traversal(T, T->root);
    printf("\n");
}

int main() {
#define ARRAY_LENGTH    20
    int keyArray[ARRAY_LENGTH] = {24, 25, 13, 35, 23, 26, 67, 47, 38, 98, 20, 19, 17, 49, 12, 21, 9, 18, 14, 15};

    rbtree *T = malloc(sizeof(rbtree));
    assert(T != NULL);

    // TODO 为什么要为nil申请节点，为NULL可以吗？
    // 不可以，因为NULL没有颜色
    T->nil = (rbtree_node *) malloc(sizeof(rbtree_node));
    assert(T->nil != NULL);

    //
    T->nil->color = BLACK;
    T->root = T->nil;

    rbtree_node *node;
    for (int i = 0; i < ARRAY_LENGTH; i++) {
        node = (rbtree_node *) malloc(sizeof(rbtree_node));
        node->key = keyArray[i];
        node->value = NULL;

        rbtree_insert(T, node);
//        rbtree_print_wrapper(T);
    }

    rbtree_traversal_wrapper(T);

    free(T->nil);
    free(T);
    return 0;
}
